package util;

import model.GameRecord;
import model.GameState;
import model.GameType;
import model.Position;
import org.apache.log4j.Logger;
import org.encog.engine.network.activation.ActivationSigmoid;
import org.encog.ml.data.basic.BasicMLData;
import org.encog.ml.data.buffer.BufferedMLDataSet;
import org.encog.ml.train.MLTrain;
import org.encog.neural.networks.BasicNetwork;
import org.encog.neural.networks.layers.BasicLayer;
import org.encog.neural.networks.training.propagation.resilient.ResilientPropagation;

import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * @author Alex Vikharev vikharev@yandex-team.ru
 *         created 28.02.12
 */
public class NeuralUtil {
    private static final Logger log = Logger.getLogger(NeuralUtil.class);

    public static void generateNeuralTrainData(final String inputFile, final String outputFile) throws IOException {
        final List<GameRecord> gameRecords = GameLoader.load(inputFile, GameType.IMPROVED);
        log.info("Loaded " + gameRecords.size() + " game records");
        Map<Position, List<Double>> scores = new HashMap<>();
        int processed = 0;
        for (GameRecord gameRecord : gameRecords) {
            final List<Position> winMoves;
            final List<Position> loseMoves;
            if (gameRecord.getResult() == GameState.FIRST_PLAYER_WIN) {
                winMoves = gameRecord.getFirstPlayerMoves();
                loseMoves = gameRecord.getSecondPlayerMoves();
            } else if (gameRecord.getResult() == GameState.SECOND_PLAYER_WIN) {
                winMoves = gameRecord.getSecondPlayerMoves();
                loseMoves = gameRecord.getFirstPlayerMoves();
            } else {
                continue;
            }
            addMoves(scores, winMoves, 1);
            addMoves(scores, loseMoves, 0);
            if (++processed % 100 == 0) {
                log.info("Processed " + processed + " game records");
            }
        }
        log.info("Processing records done. Start writing training data.");

        BufferedMLDataSet dataSet = new BufferedMLDataSet(new File(outputFile));
        dataSet.beginLoad(17 * 5, 1);
        for (Position position : scores.keySet()) {
            double score = avg(scores.get(position));
            final double[] converted = Util.BITWISE_CONVERTER.apply(position);
            final BasicMLData input = new BasicMLData(converted);
            final BasicMLData output = new BasicMLData(new double[]{score});
            dataSet.add(input, output);
        }
        dataSet.endLoad();
        log.info("Done.");
    }

    private static void addMoves(Map<Position, List<Double>> scores, List<Position> positions, double mark) {
        final List<Double> positionScores = getScores(positions.size(), 0.5, mark);
        for (int i = 0; i < positions.size(); i++) {
            final Position position = Util.getCanonical(positions.get(i));
            final Double score = positionScores.get(i);
            if (!scores.containsKey(position)) {
                scores.put(position, new ArrayList<Double>());
            }
            scores.get(position).add(score);
        }
    }

    private static double avg(List<Double> numbers) {
        double count = numbers.size();
        double sum = 0;
        for (Double number : numbers) {
            sum += number;
        }
        return sum / count;
    }

    private static List<Double> getScores(int count, double first, double last) {
        List<Double> result = new ArrayList<>(count);
        final double delta = (last - first) / (count - 1);
        for (int i = 0; i < count; i++) {
            result.add(first + i * delta);
        }
        return result;
    }

    private static void teachNetwork(String inputFile, String outputFile) {
        BasicNetwork network = new BasicNetwork();
        network.addLayer(new BasicLayer(17 * 5));
        network.addLayer(new BasicLayer(new ActivationSigmoid(), true, 17 * 5));
        network.addLayer(new BasicLayer(new ActivationSigmoid(), true, 17 * 5));
        network.addLayer(new BasicLayer(new ActivationSigmoid(), true, 1));
        network.getStructure().finalizeStructure();
        network.reset();

        BufferedMLDataSet dataSet = new BufferedMLDataSet(new File(inputFile));
        MLTrain train = new ResilientPropagation(network, dataSet);
        int epoch = 0;
        double error;
        log.info("Start teaching bot");
        do {
            train.iteration();
            ++epoch;
            error = train.getError();
            log.info("Training network, epoch: " + epoch + " error: " + error);

        } while (error > 0.01 && epoch < 1000);
        log.info("Training is over, epoch: " + epoch + " error: " + error);
        train.finishTraining();
        Util.saveClassifier(network, outputFile);
    }

    public static void main(String[] args) throws IOException {
//        generateNeuralTrainData("games10000_S", "train_network_new");
        teachNetwork("train_network_new", "network.ser");
    }
}
